import os
import unittest

import torch
from torch.backends import cudnn
from torchsummary import summary
from models.deeplab_v3_plus import DeepLabV3Plus

torch.backends.cudnn.benchmark = True

os.environ["CUDA_VISIBLE_DEVICES"] = '0'


class DeepLabV3PlusTestCase(unittest.TestCase):

    def test_summary_deeplab(self):
        self.assertTrue(torch.cuda.is_available())
        model = DeepLabV3Plus(backbone='xception')

        inputs = torch.rand(2, 3, 512, 512)
        output = model(inputs)
        print(output.size())

        # model = model.cuda()
        input_size = (3, 512, 512)
        summary(model, input_size=input_size, device='cpu')


if __name__ == '__main__':
    unittest.main()
